package eggfly.toyvm;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.bytecode.*;
import org.apache.commons.lang3.ArrayUtils;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.Stack;

public class VMMain {
    private static class Frame {
        String methodSignature;
        ArrayList<Object> localVarArray = new ArrayList<>();
        Stack<Object> operandStack = new Stack<>();

        @Override
        public String toString() {
            return methodSignature;
        }
    }

    private static class FrameStack extends Stack<Frame> {
        @Override
        public synchronized String toString() {
            return dumpStack();
        }

        String dumpStack() {
            // TODO
            return null;
        }

        Frame getCurrentFrame() {
            return this.peek();
        }
    }

    private static ThreadLocal<FrameStack> sFrameStacks = ThreadLocal.withInitial(FrameStack::new);

    public static void main(String[] args) {
        System.out.println("------------------------- VMMain ------------------------");
        try {
            innerMain();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @SuppressWarnings("SpellCheckingInspection")
    private static void innerMain() throws NotFoundException {
        ClassPool classPool = ClassPool.getDefault();
        classPool.appendClassPath("toyvm-hello/build/classes/java/main/");
        CtClass appMainClass = classPool.get("eggfly.demo.Main");
        CtMethod mainMethod = appMainClass.getDeclaredMethod("main", new CtClass[]{classPool.get("java.lang.String[]")});
        Frame newFrame = new Frame();
        newFrame.methodSignature = mainMethod.getLongName();
        // This is only a fake string array of main method's String[] param
        newFrame.localVarArray.add(new Object());
        sFrameStacks.get().push(newFrame);
        // invoke Main.main()
        invoke(classPool, appMainClass, mainMethod);
        System.err.println("ok");
    }

    @SuppressWarnings("SpellCheckingInspection")
    private static void invoke(ClassPool classPool, CtClass classOfMethod, CtMethod method) {
        if (!method.getDeclaringClass().equals(classOfMethod)) {
            throw new IllegalArgumentException("classOfMethod doesn't match method! classOfMethod=" + classOfMethod + ", method=" + method);
        }
        ConstPool classConstPool = classOfMethod.getClassFile2().getConstPool();
        MethodInfo methodInfo2 = method.getMethodInfo2();
        CodeAttribute code2 = methodInfo2.getCodeAttribute();
        try {
            System.out.println("/ --> class file path: " + classOfMethod.getURL());
        } catch (NotFoundException e) {
            e.printStackTrace();
        }
        if (code2 == null) {
            // TODO invoke native/abstract or interface method
            System.out.println("--> TODO: invoke native method: " + method.getLongName());
        } else {
            CodeIterator iterator = code2.iterator();
            int codeLength = iterator.getCodeLength();
            System.out.println("--> invoke method: " + method.getLongName() + ", code codeLength=" + codeLength);
            LocalVariableAttribute localVarsTable = (LocalVariableAttribute) code2.getAttribute(LocalVariableAttribute.tag);
            // 貌似是范型
//        LocalVariableTypeAttribute localVarTypesTable = (LocalVariableTypeAttribute) code2.getAttribute(LocalVariableTypeAttribute.tag);
            outer:
            while (iterator.hasNext()) {
                try {
                    int pos = iterator.next();
                    int opcode = iterator.byteAt(pos);
                    String opcodeName = getOpcodeName(opcode);
                    if (opcodeName == null) {
                        throw new IllegalArgumentException("opcode name not found, opcode=" + opcode);
                    }
                    System.out.println(pos + " --> opcode: " + opcodeName);
                    switch (opcode) {
                        case Opcode.SIPUSH: {
                            int shortImmediate = iterator.s16bitAt(pos + 1);
                            // push into top frame's operand stack
                            pushIntoTopFrameOperandStack(shortImmediate);
                            System.out.println(pos + ": SIPUSH: value=" + shortImmediate);
                        }
                        break;
                        case Opcode.GETSTATIC: {
                            int index = iterator.u16bitAt(pos + 1);
                            String className = classConstPool.getFieldrefClassName(index);
                            String fieldName = classConstPool.getFieldrefName(index);
                            System.out.println(pos + ": GETSTATIC: " + className + "." + fieldName);
                        }
                        break;
                        case Opcode.LDC: {
                            int index = iterator.byteAt(pos + 1);
                            Object value = classConstPool.getLdcValue(index);
                            pushIntoTopFrameOperandStack(value);
                            System.out.println(pos + ": LDC: " + value);
                        }
                        break;
                        case Opcode.LDC2_W: {
                            int index = iterator.u16bitAt(pos + 1);
                            Object value = classConstPool.getLdcValue(index);
                            pushIntoTopFrameOperandStack(value);
                            System.out.println(pos + ": LDC2_W: " + value);
                        }
                        break;
                        case Opcode.INVOKESTATIC: {
                            handleInvokeRelatedOp("INVOKESTATIC", classPool, classConstPool, iterator, pos);
                        }
                        break;
                        case Opcode.INVOKEVIRTUAL: {
                            // threadStack.push(thisVar);
                            handleInvokeRelatedOp("INVOKEVIRTUAL", classPool, classConstPool, iterator, pos);
                        }
                        break;
                        case Opcode.INVOKEDYNAMIC: {
                            handleInvokeRelatedOp("INVOKEDYNAMIC", classPool, classConstPool, iterator, pos);
                        }
                        break;
                        case Opcode.INVOKEINTERFACE: {
                            handleInvokeRelatedOp("INVOKEINTERFACE", classPool, classConstPool, iterator, pos);
                        }
                        break;
                        case Opcode.INVOKESPECIAL: {
                            handleInvokeRelatedOp("INVOKESPECIAL", classPool, classConstPool, iterator, pos);
                        }
                        break;
                        case Opcode.POP: {
                            popFromTopFrameOperandStack();
                            System.out.println(pos + ": POP");
                        }
                        break;
                        case Opcode.POP2: {
                            // same with "POP" because long type costs two slots
                            popFromTopFrameOperandStack();
                            System.out.println(pos + ": POP2");
                        }
                        break;
                        case Opcode.RETURN: {
                            // clear top frame
                            sFrameStacks.get().pop();
                            System.out.println(pos + ": RETURN");
                            break outer;
                        }
                        case Opcode.LRETURN: {
                            long value = popFromTopFrameOperandStackWithTypeCheck(Long.class, opcodeName);
                            // clear top frame
                            sFrameStacks.get().pop();
                            sFrameStacks.get().getCurrentFrame().operandStack.push(value);
                            System.out.println(pos + ": LRETURN");
                            break outer;
                        }
                        case Opcode.ICONST_M1: {
                            pushIntoTopFrameOperandStack(-1);
                            System.out.println(pos + ": ICONST_M1");
                        }
                        break;
                        case Opcode.LCONST_0: {
                            pushIntoTopFrameOperandStack(0L);
                            System.out.println(pos + ": LCONST_0");
                        }
                        break;
                        case Opcode.LCMP: {
                            long value1 = popFromTopFrameOperandStackWithTypeCheck(Long.class, opcodeName);
                            long value2 = popFromTopFrameOperandStackWithTypeCheck(Long.class, opcodeName);
                            final int result = Long.compare(value1, value2);
                            pushIntoTopFrameOperandStack(result);
                            System.out.println(pos + ": LCMP");
                        }
                        break;
                        case Opcode.IFNE: {
                            int value = popFromTopFrameOperandStackWithTypeCheck(Integer.class, opcodeName);
                            if (value != 0) {
                                int relativeOffset = iterator.s16bitAt(pos + 1);
                                pos += relativeOffset;
                                iterator.move(pos);
                            }
                            System.out.println(pos + ": IFNE");
                        }
                        break;
                        case Opcode.ILOAD_0: {
                            Object localVar0 = sFrameStacks.get().getCurrentFrame().localVarArray.get(0);
                            if (!(localVar0 instanceof Integer)) {
                                throw new IllegalStateException("The localVar0 is not a integer, " +
                                        "but opcode is ILOAD_0, maybe bytecode error or vm internal error!");
                            }
                            pushIntoTopFrameOperandStack(localVar0);
                            System.out.println(pos + ": ILOAD_0");
                        }
                        break;
                        case Opcode.LLOAD_1: {
                            Object localVar1 = sFrameStacks.get().getCurrentFrame().localVarArray.get(1);
                            if (!(localVar1 instanceof Long)) {
                                throw new IllegalStateException("The localVar1 is not a long, " +
                                        "but opcode is LLOAD_1, maybe bytecode error or vm internal error!");
                            }
                            pushIntoTopFrameOperandStack(localVar1);
                            System.out.println(pos + ": LLOAD_1");
                        }
                        break;
                        case Opcode.LSTORE_1: {
                            handleStoreByTypeAndIndex(Long.class, 1, opcodeName);
                            System.out.println(pos + ": LSTORE_1");
                        }
                        break;
                        case Opcode.I2L: {
                            int value = popFromTopFrameOperandStackWithTypeCheck(Integer.class, opcodeName);
                            pushIntoTopFrameOperandStack((long) value);
                            System.out.println(pos + ": I2L");
                        }
                        break;
                        case Opcode.LADD: {
                            Object oprand1 = popFromTopFrameOperandStack();
                            Object oprand2 = popFromTopFrameOperandStack();
                            if (!(oprand1 instanceof Long)) {
                                throw new IllegalStateException("The top1 value from operand stack is not a long, " +
                                        "but opcode is LADD, maybe bytecode error or vm internal error!");
                            }
                            if (!(oprand2 instanceof Long)) {
                                throw new IllegalStateException("The top2 value from operand stack is not a long, " +
                                        "but opcode is LADD, maybe bytecode error or vm internal error!");
                            }
                            long longValue1 = (long) oprand1;
                            long longValue2 = (long) oprand2;
                            pushIntoTopFrameOperandStack(longValue1 + longValue2);
                            System.out.println(pos + ": LADD");
                        }
                        break;
                        case Opcode.NOP: {
                            System.out.println(pos + ": NOP");
                        }
                        break;
                        default: {
                            System.err.println("This opcode is known, but not handled by ToyVM yet, opcode = " + opcodeName);
                        }
                        break;
                    }
                } catch (BadBytecode badBytecode) {
                    badBytecode.printStackTrace();
                }
            }
        }
    }

    private static <T> void handleStoreByTypeAndIndex(Class<T> type, int index, String opcodeName) {
        T value = popFromTopFrameOperandStackWithTypeCheck(type, opcodeName);
        ArrayList<Object> localVarArray = sFrameStacks.get().getCurrentFrame().localVarArray;
        if (localVarArray.size() > index) {
            localVarArray.set(index, value);
        } else if (localVarArray.size() == index) {
            localVarArray.add(value);
        } else {
            throw new IllegalStateException("Current opcode is " + opcodeName
                    + ", but localVariableArray.size() = " + localVarArray.size() + ", cannot handle STORE opcodes.");
        }
    }

    private static <T> T popFromTopFrameOperandStackWithTypeCheck(Class<T> type, String opcodeName) {
        Object value = popFromTopFrameOperandStack();
        // TODO: NPE?
        if (value.getClass() != type) {
            throw new IllegalStateException("The top value's type from operand stack is not " + type + ", " +
                    "but opcode is " + opcodeName + ", maybe bytecode error or vm internal error!");
        }
        //noinspection unchecked
        return (T) value;
    }

    private static void pushIntoTopFrameOperandStack(Object value) {
        sFrameStacks.get().getCurrentFrame().operandStack.push(value);
    }

    private static Object popFromTopFrameOperandStack() {
        return sFrameStacks.get().getCurrentFrame().operandStack.pop();
    }

    private static String getOpcodeName(int opcode) {
        Field[] fields = Opcode.class.getDeclaredFields();
        for (Field field : fields) {
            try {
                if (field.getType() == int.class && field.get(null).equals(opcode)) {
                    return field.getName();
                }
            } catch (IllegalAccessException e) {
                e.printStackTrace();
            }
        }
        return null;
    }

    private static void handleInvokeRelatedOp(String invokeTypeForLog, ClassPool classPool, ConstPool classConstPool, CodeIterator iterator, int pos) {
        int index = iterator.u16bitAt(pos + 1);
        String className = classConstPool.getMethodrefClassName(index);
        String methodName = classConstPool.getMethodrefName(index);
        String methodType = classConstPool.getMethodrefType(index);
        try {
            CtClass methodClass = classPool.get(className);
            CtMethod[] overrideMethods = methodClass.getDeclaredMethods(methodName);
            CtMethod foundMethod = null;
            for (CtMethod overrideMethod : overrideMethods) {
                if (overrideMethod.getSignature().equals(methodType)) {
                    // TODO: 暂时不考虑范型? overrideMethod.getGenericSignature()
                    foundMethod = overrideMethod;
                    break;
                }
            }
            if (foundMethod != null) {
                if ("java.io.PrintStream".equals(className) && "println".equals(methodName)) {
                    Object objToPrint = popFromTopFrameOperandStack();
                    System.out.println("Intercept calling java.io.PrintStream.println()");
                    System.out.println("--- start of print ---");
                    System.out.println(objToPrint);
                    System.out.println("--- end of print ---");
                } else {
                    CtClass[] paramTypes = foundMethod.getParameterTypes();
                    Frame currentFrame = sFrameStacks.get().getCurrentFrame();
                    Frame newFrame = new Frame();
                    newFrame.methodSignature = foundMethod.getLongName();
                    sFrameStacks.get().push(newFrame);
                    StringBuilder builder = new StringBuilder("* params: ");
                    Object[] reversed = new Object[paramTypes.length];
                    for (int i = 0; i < paramTypes.length; i++) {
                        Object paramValue = currentFrame.operandStack.pop();
                        reversed[i] = paramValue;
                    }
                    ArrayUtils.reverse(reversed);
                    for (int i = 0; i < reversed.length; i++) {
                        Object value = reversed[i];
                        newFrame.localVarArray.add(value);
                        builder.append('[').append(i).append(']').append("=").append(value).append(", ");
                    }
                    System.out.println(builder.toString());
                    invoke(classPool, methodClass, foundMethod);
                }
            }
        } catch (NotFoundException e) {
            e.printStackTrace();
        }
        System.out.println(pos + ": " + invokeTypeForLog + ": " + className + "." + methodName);
    }

    private static void dump() {
//        try {
//            javassist.tools.Dump.main(new String[]{appMainClass.getURL().getPath()});
//        } catch (Exception e) {
//            e.printStackTrace();
//        }
    }
}
